#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
瞬变源数据展示Flask应用
用于展示瞬变源分类结果，包括图像和属性信息
"""

import os
import json
from flask import Flask, render_template, request, jsonify, send_from_directory

app = Flask(__name__, template_folder='templates', static_folder='static')

# 数据目录
DATASETS_DIR = 'data/datasets'
CLASSIFY_TYPE_DIR = 'funs/config_data'
OT_CLASSIFY_SERVICE_DATA_DIR = 'ot_classify_service/data'

def get_datasets():
    """获取所有样本集"""
    datasets = []
    
    # 获取 data/datasets 目录下的样本集
    if os.path.exists(DATASETS_DIR):
        for item in os.listdir(DATASETS_DIR):
            item_path = os.path.join(DATASETS_DIR, item)
            if os.path.isdir(item_path):
                datasets.append({
                    'name': item,
                    'source': 'data/datasets'
                })
    
    # 获取 ot_classify_service/data 目录下的样本集
    if os.path.exists(OT_CLASSIFY_SERVICE_DATA_DIR):
        for item in os.listdir(OT_CLASSIFY_SERVICE_DATA_DIR):
            item_path = os.path.join(OT_CLASSIFY_SERVICE_DATA_DIR, item)
            if os.path.isdir(item_path):
                datasets.append({
                    'name': item,
                    'source': 'ot_classify_service/data'
                })
    
    return datasets

def load_transient_data(dataset_name, source_dir='data/datasets'):
    """加载指定样本集的瞬变源数据"""
    data_file = os.path.join(source_dir, dataset_name, 'transient_classify_rst.json')
    if os.path.exists(data_file):
        with open(data_file, 'r', encoding='utf-8') as f:
            return json.load(f)
    return {'objects': [], 'obj_number': 0}

def save_transient_data(dataset_name, data, source_dir='data/datasets'):
    """保存指定样本集的瞬变源数据"""
    data_file = os.path.join(source_dir, dataset_name, 'transient_classify_rst.json')
    with open(data_file, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=2)

def load_transient_types():
    """加载瞬变源分类类型"""
    type_file = os.path.join(CLASSIFY_TYPE_DIR, 'transient_type.json')
    if os.path.exists(type_file):
        with open(type_file, 'r', encoding='utf-8') as f:
            return json.load(f)
    return {
        "is_source": "有目标",
        "is_hot_pixel": "热像素",
        "is_bright_star_residual": "亮星残差",
        "is_transient": "瞬变源",
        "is_dark_pixel": "暗像素",
        "other": "未知类型"
    }

@app.route('/')
def index():
    """首页：展示样本集列表"""
    datasets = get_datasets()
    return render_template('index.html', datasets=datasets)

@app.route('/dataset/<path:dataset_path>')
def dataset_detail(dataset_path):
    """样本集详情页：展示该样本集的瞬变源列表，支持分页和分类筛选"""
    # 解析数据集路径，格式为 "source_dir/dataset_name"
    path_parts = dataset_path.split('/', 1)
    if len(path_parts) != 2:
        return "无效的数据集路径", 400
    
    source_dir, dataset_name = path_parts
    
    # 获取分页参数
    page = request.args.get('page', 1, type=int)
    per_page = request.args.get('per_page', 30, type=int)
    filter_classification = request.args.get('filter', None)
    
    # 限制每页显示数量选项
    per_page_options = [30, 50, 100]
    if per_page not in per_page_options:
        per_page = 30
    
    # 确定实际的数据目录
    if source_dir == 'data':
        actual_source_dir = 'data/datasets'
    elif source_dir == 'ot_classify_service':
        actual_source_dir = 'ot_classify_service/data'
    else:
        return "无效的数据源", 400
    
    # 加载数据
    data = load_transient_data(dataset_name, actual_source_dir)
    objects = data.get('objects', [])
    
    # 如果有筛选条件，只保留指定分类的目标
    if filter_classification:
        objects = [obj for obj in objects if obj.get('classification') == filter_classification]
    
    # 计算分页信息
    total_count = len(objects)
    total_pages = (total_count + per_page - 1) // per_page
    page = max(1, min(page, total_pages))
    start_index = (page - 1) * per_page
    end_index = start_index + per_page
    paginated_objects = objects[start_index:end_index]
    
    # 统计各类别的数量（基于所有目标，不考虑筛选）
    all_objects = data.get('objects', [])
    classification_count = {}
    for obj in all_objects:
        classification = obj.get('classification', '未知类型')
        classification_count[classification] = classification_count.get(classification, 0) + 1
    
    # 加载瞬变源分类类型
    transient_types = load_transient_types()
    
    return render_template('dataset_detail.html', 
                         dataset_name=dataset_name,
                         objects=paginated_objects,
                         total_count=total_count,
                         classification_count=classification_count,
                         page=page,
                         per_page=per_page,
                         per_page_options=per_page_options,
                         total_pages=total_pages,
                         transient_types=transient_types)

@app.route('/dataset/<path:dataset_path>/transient/<name>')
def transient_detail(dataset_path, name):
    """瞬变源详情页"""
    # 解析数据集路径，格式为 "source_dir/dataset_name"
    path_parts = dataset_path.split('/', 1)
    if len(path_parts) != 2:
        return "无效的数据集路径", 400
    
    source_dir, dataset_name = path_parts
    
    # 确定实际的数据目录
    if source_dir == 'data':
        actual_source_dir = 'data/datasets'
    elif source_dir == 'ot_classify_service':
        actual_source_dir = 'ot_classify_service/data'
    else:
        return "无效的数据源", 400
    
    data = load_transient_data(dataset_name, actual_source_dir)
    objects = data.get('objects', [])
    
    # 查找指定的瞬变源
    transient = None
    for obj in objects:
        if obj.get('name') == name:
            transient = obj
            break
    
    if transient is None:
        return "瞬变源未找到", 404
    
    # 获取分类判断依据信息
    from funs.transient_classifier import TransientClassify
    import types
    
    # 创建一个模拟的Transient对象，只包含需要的属性
    class MockTransient:
        def __init__(self, data):
            # 从数据中提取所有需要的属性
            self.peak_flux_obs = data['obs']['peak_flux']
            self.obs_rms = data['obs']['rms']
            self.peak_flux_tpl = data['tpl']['peak_flux']
            self.tpl_rms = data['tpl']['rms']
            self.sharpness_obs = data['obs']['sharpness']
            self.sharpness_res = data['res']['sharpness']
            self.fwhm_obs = data['obs']['fwhm']
            self.fwhm_res = data['res']['fwhm']
            self.snr_obs = data['obs']['snr']
            self.peak_to_3x3_obs = data['obs']['peak_to_3x3']
            self.flux3_to_5x5_obs = data['obs']['flux3_to_5x5']
    
    transient_obj = MockTransient(transient)
    classifier = TransientClassify(None, None, None, [])  # 只需要使用其方法
    classification_criteria = classifier.get_classification_criteria(transient_obj)
    
    # 加载瞬变源分类类型
    transient_types = load_transient_types()
    
    return render_template('detail.html', dataset_name=dataset_name, transient=transient, classification_criteria=classification_criteria, transient_types=transient_types)

@app.route('/save_classification', methods=['POST'])
def save_classification():
    """保存人工分类结果"""
    data = request.get_json()
    dataset_path = data.get('dataset_path')  # 使用路径而不是名称
    transient_name = data.get('transient_name')
    manual_classify_result = data.get('manual_classify_result')
    
    # 解析数据集路径
    path_parts = dataset_path.split('/', 1)
    if len(path_parts) != 2:
        return jsonify({'status': 'error', 'message': '无效的数据集路径'}), 400
    
    source_dir, dataset_name = path_parts
    
    # 确定实际的数据目录
    if source_dir == 'data':
        actual_source_dir = 'data/datasets'
    elif source_dir == 'ot_classify_service':
        actual_source_dir = 'ot_classify_service/data'
    else:
        return jsonify({'status': 'error', 'message': '无效的数据源'}), 400
    
    # 加载数据并更新
    all_data = load_transient_data(dataset_name, actual_source_dir)
    objects = all_data.get('objects', [])
    
    for obj in objects:
        if obj.get('name') == transient_name:
            obj['manual_classify_result'] = manual_classify_result
            break
    
    # 保存更新后的数据
    save_transient_data(dataset_name, all_data, actual_source_dir)
    
    return jsonify({'status': 'success'})

@app.route('/dataset/<path:dataset_path>/images/<path:filename>')
def serve_image(dataset_path, filename):
    """提供图像文件服务"""
    # 解析数据集路径，格式为 "source_dir/dataset_name"
    path_parts = dataset_path.split('/', 1)
    if len(path_parts) != 2:
        return "无效的数据集路径", 400
    
    source_dir, dataset_name = path_parts
    
    # 确定实际的数据目录
    if source_dir == 'data':
        actual_source_dir = 'data/datasets'
    elif source_dir == 'ot_classify_service':
        actual_source_dir = 'ot_classify_service/data'
    else:
        return "无效的数据源", 400
    
    image_dir = os.path.join(actual_source_dir, dataset_name, 'images')
    return send_from_directory(image_dir, filename)

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5001)